#!/bin/bash
# Copyright 2021 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# Create the TensorFlow Decision Forests pip package.
# This command uses the compiled artifacts generated by test_bazel.sh.
#
# Requirements:
#   pyenv
#     See https://github.com/pyenv/pyenv-installer
#
#   manylinux compatible docker
#     sudo docker pull tensorflow/tensorflow:custom-op-ubuntu16
#     sudo docker run -it -v ${PWD}:/working_dir -w /working_dir  tensorflow/tensorflow:custom-op-ubuntu16
#
#   auditwheel
#     python -m pip install install auditwheel -U
#     + in auditwheel "policy.json" file, add "libtensorflow_framework.so.2" next to each "libresolv.so.2" entry.
#     See https://github.com/tensorflow/tensorflow/issues/31807
#     e.g. /home/${USER}/.local/lib/python3.9/site-packages/auditwheel/policy/policy.json
#
set -x
set -e

# Temporary directory used to assemble the package.
SRCPK="$(pwd)/tmp_package"

# Check the correct location of the current directory.
if [ ! -d "bazel-bin" ]; then
  echo "This script should be run from the root directory of TensorFlow Decision Forests (i.e. the directory containing the LICENSE file) of a compiled Bazel export (i.e. containing a bazel-bin directory)"
  exit 1
fi

function install_dependencies() {
  PYTHON="$1"; shift
  ${PYTHON} -m pip install setuptools -U
  ${PYTHON} -m pip install build -U
  ${PYTHON} -m pip install virtualenv -U
  # TODO(janpf): Automatically edit the policy.json file with the
  # libtensorflow_framework.so.2 entry.
  ${PYTHON} -m pip install auditwheel -U
}

function assemble_files() {
  rm -fr ${SRCPK}
  mkdir -p ${SRCPK}
  cp -R tensorflow_decision_forests LICENSE configure/setup.py configure/MANIFEST.in README.md ${SRCPK}

  # TFDF's wrappers and .so.
  SRCBIN="bazel-bin/tensorflow_decision_forests"
  cp ${SRCBIN}/tensorflow/ops/inference/inference.so ${SRCPK}/tensorflow_decision_forests/tensorflow/ops/inference/
  cp ${SRCBIN}/tensorflow/ops/inference/op.py ${SRCPK}/tensorflow_decision_forests/tensorflow/ops/inference/
  cp ${SRCBIN}/tensorflow/ops/training/training.so ${SRCPK}/tensorflow_decision_forests/tensorflow/ops/training/
  cp ${SRCBIN}/tensorflow/ops/training/op.py ${SRCPK}/tensorflow_decision_forests/tensorflow/ops/training/
  cp ${SRCBIN}/keras/wrappers.py ${SRCPK}/tensorflow_decision_forests/keras/

  # YDF's proto wrappers.
  YDFSRCBIN="bazel-bin/external/ydf/yggdrasil_decision_forests"
  mkdir -p ${SRCPK}/yggdrasil_decision_forests
  pushd ${YDFSRCBIN}
  find -name \*.py -exec cp --parents -prv {} ${SRCPK}/yggdrasil_decision_forests \;
  popd

  # Add __init__.py to all exported Yggdrqasil sub-directories.
  find ${SRCPK}/yggdrasil_decision_forests -type d -exec touch {}/__init__.py \;
}

function build_package() {
  PYTHON="$1"; shift

  pushd ${SRCPK}
  $PYTHON -m build
  popd

  cp -R ${SRCPK}/dist .
}

function test_package() {
  PYTHON="$1"; shift
  PACKAGE="$1"; shift

  pip3 install --upgrade pip
  pip3 install dist/tensorflow_decision_forests-*-cp${PACKAGE}-cp${PACKAGE}*-linux_x86_64.whl --force-reinstall
  pip3 list
  pip3 show tensorflow_decision_forests -f
  $PYTHON examples/minimal.py
}

function build_package_pyenv() {
  VERSION="$1"; shift
  PACKAGE="$1"; shift

  ENVNAME=env_${VERSION}
  set +e
  pyenv virtualenv ${VERSION} ${ENVNAME}
  set -e
  pyenv activate ${ENVNAME}

  install_dependencies python3

  build_package python3

  pyenv deactivate

  auditwheel show dist/tensorflow_decision_forests-*-cp${PACKAGE}-cp${PACKAGE}*-linux_x86_64.whl
}

function test_package_pyenv() {
  VERSION="$1"; shift
  PACKAGE="$1"; shift

  ENVNAME=env_${VERSION}
  set +e
  pyenv virtualenv ${VERSION} ${ENVNAME}
  set -e
  pyenv activate ${ENVNAME}

  test_package python3 ${PACKAGE}

  pyenv deactivate
}

function e2e_env() {
  VERSION="$1"; shift
  PACKAGE="$1"; shift

  pyenv install ${VERSION} -s

  build_package_pyenv ${VERSION} ${PACKAGE}
  test_package_pyenv ${VERSION} ${PACKAGE}

  auditwheel repair --plat manylinux2010_x86_64 -w dist dist/tensorflow_decision_forests-*-cp${PACKAGE}-cp${PACKAGE}*-linux_x86_64.whl
}

eval "$(pyenv init -)"

assemble_files

if [ "$1" = "simple" ]; then
  install_dependencies python3.9
  build_package python3.9 39
  test_package python3.9 39
else
  e2e_env 3.6.8 36
  e2e_env 3.7.7 37
  e2e_env 3.8.7 38
  e2e_env 3.9.2 39
fi
